﻿/*
 * Description: 显示分辨率设置修改
 */

#include <Windows.h>
#include <iostream>
#include <string>
#include <tchar.h>

namespace {
    const int Message_DELAY = 5000;
    const int RESOLUTION_ARGS_NUM = 3;
    const int MAX_RETRY_TIMES = 5;
    const std::wstring SVGADEVICEID = L"PCI\\VEN_15AD&DEV_0405";

    std::wstring GetDeviceName()
    {
        std::wstring displayDeviceName = L"";
        DISPLAY_DEVICEW adapter = { 0 };
        adapter.cb = sizeof(DISPLAY_DEVICE);
        DWORD adapterIndex = 0;
        while (EnumDisplayDevicesW(NULL, adapterIndex, &adapter, 0)) {
            if (!_wcsnicmp(adapter.DeviceID, SVGADEVICEID.c_str(), SVGADEVICEID.size())) {
                displayDeviceName = adapter.DeviceName;
                break;
            }
            adapter = { 0 };
            adapter.cb = sizeof(DISPLAY_DEVICE);
            ++adapterIndex;
        }
        return displayDeviceName;
    }

    int GetModeInfoIndex(UINT pathArrayElements, DISPLAYCONFIG_PATH_INFO* displayPathInfo, UINT modeInfoArrayElements,
                         DISPLAYCONFIG_MODE_INFO* displayModeInfo)
    {
        std::wstring displayDeviceName = GetDeviceName();
        if (QueryDisplayConfig(QDC_ALL_PATHS, &pathArrayElements, displayPathInfo, &modeInfoArrayElements,
            displayModeInfo, NULL) != ERROR_SUCCESS) {
                return -1;
        }
        for (UINT pathIdx = 0; pathIdx < pathArrayElements; ++pathIdx) {
            DISPLAYCONFIG_SOURCE_DEVICE_NAME sourceName = {};
            sourceName.header.type = DISPLAYCONFIG_DEVICE_INFO_GET_SOURCE_NAME;
            sourceName.header.size = sizeof(sourceName);
            sourceName.header.adapterId = displayPathInfo[pathIdx].sourceInfo.adapterId;
            sourceName.header.id = displayPathInfo[pathIdx].sourceInfo.id;
            if (DisplayConfigGetDeviceInfo(&sourceName.header) != ERROR_SUCCESS) {
                continue;
            }
            if (wcscmp(displayDeviceName.c_str(), sourceName.viewGdiDeviceName) == 0) {
                return pathIdx;
            }
        }

        return -1;
    }

    bool ChangeDisplay(const DWORD targetWidth, const DWORD targetHeight)
    {
        UINT pathArrayElements = 0;
        UINT modeInfoArrayElements = 0;
        if (GetDisplayConfigBufferSizes(QDC_ALL_PATHS, &pathArrayElements, &modeInfoArrayElements) != ERROR_SUCCESS) {
            return false;
        }
        DISPLAYCONFIG_PATH_INFO* displayPathInfo = new DISPLAYCONFIG_PATH_INFO[pathArrayElements];
        DISPLAYCONFIG_MODE_INFO* displayModeInfo = new DISPLAYCONFIG_MODE_INFO[modeInfoArrayElements];
        SecureZeroMemory(displayPathInfo, sizeof(DISPLAYCONFIG_PATH_INFO) * pathArrayElements);
        SecureZeroMemory(displayModeInfo, sizeof(DISPLAYCONFIG_MODE_INFO) * modeInfoArrayElements);
        int pathArrayIndex = GetModeInfoIndex(pathArrayElements, displayPathInfo, modeInfoArrayElements,
            displayModeInfo);
        if (pathArrayIndex == -1) {
            delete[] displayPathInfo;
            delete[] displayModeInfo;
            displayPathInfo = NULL;
            displayModeInfo = NULL;
            return false;
        }
        UINT modeIndex = displayPathInfo[pathArrayIndex].sourceInfo.modeInfoIdx;
        DISPLAYCONFIG_MODE_INFO modeInfo = displayModeInfo[modeIndex];
        
        displayModeInfo[modeIndex].sourceMode.width = targetWidth;
        displayModeInfo[modeIndex].sourceMode.height = targetHeight;

        SetDisplayConfig(pathArrayElements, displayPathInfo, modeInfoArrayElements, displayModeInfo, SDC_APPLY |
            SDC_USE_SUPPLIED_DISPLAY_CONFIG | SDC_ALLOW_CHANGES | SDC_SAVE_TO_DATABASE);
        
        for (int i = 0; i < MAX_RETRY_TIMES; i++) {
            if (QueryDisplayConfig(QDC_ALL_PATHS, &pathArrayElements, displayPathInfo, &modeInfoArrayElements,
                displayModeInfo, NULL) == ERROR_SUCCESS) {
                if (displayModeInfo[modeIndex].sourceMode.width != targetWidth
                    || displayModeInfo[modeIndex].sourceMode.height != targetHeight) {
                    displayModeInfo[modeIndex].sourceMode.width = targetWidth;
                    displayModeInfo[modeIndex].sourceMode.height = targetHeight;
                    SetDisplayConfig(pathArrayElements, displayPathInfo, modeInfoArrayElements, displayModeInfo,
                        SDC_APPLY | SDC_USE_SUPPLIED_DISPLAY_CONFIG | SDC_ALLOW_CHANGES | SDC_SAVE_TO_DATABASE);
                } else {
                    break;
                }
            }
            Sleep(500);  // 等待500ms后重试
        }

        delete[] displayPathInfo;
        delete[] displayModeInfo;
        displayPathInfo = NULL;
        displayModeInfo = NULL;
        return true;
    }

    bool IsResolutionSupported(uint32_t targetWidth, uint32_t targetHeight)
    {
        DISPLAY_DEVICE adapter = { 0 };
        adapter.cb = sizeof(DISPLAY_DEVICE);
        DWORD adapterIndex = 0;
        while (EnumDisplayDevices(nullptr, adapterIndex, &adapter, 0)) {
            DEVMODE mode = { 0 };
            mode.dmSize = sizeof(DEVMODE);
            DWORD modeIndex = 0;
            while (EnumDisplaySettings(adapter.DeviceName, modeIndex, &mode)) {
                if (mode.dmPelsWidth != targetWidth || mode.dmPelsHeight != targetHeight) {
                    mode = { 0 };
                    mode.dmSize = sizeof(DEVMODE);
                    ++modeIndex;
                    continue;
                }
                return true;
            }
            adapter = { 0 };
            adapter.cb = sizeof(DISPLAY_DEVICE);
            ++adapterIndex;
        }
        return false;
    }

	int InitResolution(uint32_t targetWidth, uint32_t targetHeight)
	{
        for (int i = 0; i < MAX_RETRY_TIMES; i++) {
            if (IsResolutionSupported(targetWidth, targetHeight) && ChangeDisplay(targetWidth, targetHeight)) {
                return 0;
            }
            Sleep(500);  // 等待500ms后重试
        }
        return -1;
	}

        //Functions & other definitions required-->
    typedef int (__stdcall *MSGBOXAAPI)(IN HWND hWnd,
        IN LPCSTR lpText, IN LPCSTR lpCaption,
        IN UINT uType, IN WORD wLanguageId, IN DWORD dwMilliseconds);
    typedef int (__stdcall *MSGBOXWAPI)(IN HWND hWnd,
        IN LPCWSTR lpText, IN LPCWSTR lpCaption,
        IN UINT uType, IN WORD wLanguageId, IN DWORD dwMilliseconds);
    
    int MessageBoxTimeoutA(IN HWND hWnd, IN LPCSTR lpText,
        IN LPCSTR lpCaption, IN UINT uType,
        IN WORD wLanguageId, IN DWORD dwMilliseconds);
    int MessageBoxTimeoutW(IN HWND hWnd, IN LPCWSTR lpText,
        IN LPCWSTR lpCaption, IN UINT uType,
        IN WORD wLanguageId, IN DWORD dwMilliseconds);
    
    #ifdef UNICODE
        #define MessageBoxTimeout MessageBoxTimeoutW
    #else
        #define MessageBoxTimeout MessageBoxTimeoutA
    #endif
    
    #define MB_TIMEDOUT 32000
    
    int MessageBoxTimeoutA(HWND hWnd, LPCSTR lpText,
        LPCSTR lpCaption, UINT uType, WORD wLanguageId,
        DWORD dwMilliseconds)
    {
        static MSGBOXAAPI MsgBoxTOA = nullptr;
    
        if (!MsgBoxTOA) {
            HMODULE hUser32 = GetModuleHandle(_T("user32.dll"));
            if (hUser32) {
                MsgBoxTOA = (MSGBOXAAPI)GetProcAddress(hUser32, "MessageBoxTimeoutA");
                //fall through to 'if (MsgBoxTOA)...'
            } else {
                //stuff happened, add code to handle it here
                //(possibly just call MessageBox())
                return 0;
            }
        }
    
        if (MsgBoxTOA) {
            return MsgBoxTOA(hWnd, lpText, lpCaption, uType, wLanguageId, dwMilliseconds);
        }
    
        return 0;
    }
    
    int MessageBoxTimeoutW(HWND hWnd, LPCWSTR lpText,
        LPCWSTR lpCaption, UINT uType, WORD wLanguageId, DWORD dwMilliseconds)
    {
        static MSGBOXWAPI MsgBoxTOW = nullptr;
    
        if (!MsgBoxTOW) {
            HMODULE hUser32 = GetModuleHandle(_T("user32.dll"));
            if (hUser32) {
                MsgBoxTOW = (MSGBOXWAPI)GetProcAddress(hUser32, "MessageBoxTimeoutW");
                //fall through to 'if (MsgBoxTOW)...'
            } else {
                //stuff happened, add code to handle it here
                //(possibly just call MessageBox())
                return 0;
            }
        }
    
        if (MsgBoxTOW) {
            return MsgBoxTOW(hWnd, lpText, lpCaption, uType, wLanguageId, dwMilliseconds);
        }
    
        return 0;
    }
}

int main(int argc, char* argv[])
{
    if (argc == 2 && strcmp(argv[1], "update") == 0) { // 2个入参时判断执行内容
        HMODULE hUser32 = LoadLibrary(_T("user32.dll"));
        if (hUser32) {
            int iRet = 0;
            UINT uiFlags = MB_OK | MB_DEFBUTTON1 | MB_APPLMODAL;
            std::string processName = "explorer.exe";
            iRet = MessageBoxTimeout(NULL, _T("即将更新虚拟机驱动，有可能造成数秒短暂黑屏\n自动恢复后可正常使用。"),
                _T("虚拟机管理助手"), uiFlags, 0, Message_DELAY);
            FreeLibrary(hUser32);
        }
        return 0;
    } else if (argc == RESOLUTION_ARGS_NUM) {
        int targetWidth = atoi(argv[1]);
        int targetHeight = atoi(argv[2]);
        if (targetWidth > 0 && targetHeight > 0) {
            return InitResolution(static_cast<uint32_t>(targetWidth), static_cast<uint32_t>(targetHeight));
        } else {
            return -1;
        }
    } else {
        return -1;
    }
}